package com.starhub.application.collection.common;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.collection.dto.DataManipulationDto;
import com.starhub.application.interfaces.model.ModelConfigService;
import com.starhub.application.model.enums.VectorModelTypeEnum;
import com.starhub.application.model.util.DimensionUtil;
import com.starhub.application.rag.dto.EmbeddingDto;
import com.starhub.application.rag.service.EmbeddingService;
import com.starhub.application.rag.service.SegmentService;
import com.starhub.common.bean.model.ModelConfig;
import com.starhub.utils.Util;

import lombok.extern.slf4j.Slf4j;

/**
 * 向量化处理类
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-09
 */
@Slf4j
@Component
public class EmbeddingHelper {

    @Autowired
    private ModelConfigService modelConfigService;

    @Autowired
    private EmbeddingService ragService;

    @Autowired
    private SegmentService segmentService;

    /**
     * 进行向量化处理
     * 
     * @param dataManipulationDto
     */
    public void embeddingData(DataManipulationDto dataManipulationDto) {
        EmbeddingDto embeddingDto = buildEmbeddingDto(dataManipulationDto);
        // 进行向量化处理
        try {
            embedding(embeddingDto);
        } catch (Exception e) {
            throw new RuntimeException("向量化失败: " + e.getMessage(), e);
        }
    }

    /**
     * 更新向量数据
     * @param dataManipulationDto
     * @return
     */
    public boolean embeddingUpdateData(DataManipulationDto dataManipulationDto) {
        EmbeddingDto embeddingDto = buildEmbeddingDto(dataManipulationDto);
        try {
            // 删除向量数据
            ragService.deleteEmbedding(embeddingDto);
            // 删除段落数据
            segmentService.deleteByItemId(Util.getLongValue(embeddingDto.getItemId()));
            // 重新向量化
            embedding(embeddingDto);
            return true;
        } catch (Exception e) {
            throw new RuntimeException("更新向量数据失败: " + e.getMessage(), e);
        }
    }

    /**
     * 删除向量数据
     * @param embeddingDto
     * @return 是否删除成功
     */
    public boolean embeddingDeleteData(DataManipulationDto dataManipulationDto) {
        EmbeddingDto embeddingDto = buildEmbeddingDto(dataManipulationDto);
        // 删除段落数据
        segmentService.deleteByItemId(Util.getLongValue(embeddingDto.getItemId()));
        return ragService.deleteEmbedding(embeddingDto);
    }

    /**
     * 向量化
     * @param embeddingDto
     * @throws Exception
     */
    private void embedding(EmbeddingDto embeddingDto) throws Exception {
        ragService.embedding(embeddingDto);
    }

    /**
     * 构建向量化处理DTO
     * 
     * @param dataManipulationDto
     * @return
     */
    private EmbeddingDto buildEmbeddingDto(DataManipulationDto dataManipulationDto) {
        ModelConfig config = modelConfigService.getModel(dataManipulationDto.getKnowledgeBase().getVectorModel());

        EmbeddingDto embeddingDto = EmbeddingDto.builder()
                .itemId(Util.null2String(dataManipulationDto.getItemId()))
                .collectionId(Util.null2String(dataManipulationDto.getCollectionId()))
                .baseId(Util.null2String(dataManipulationDto.getBaseId()))
                .content(dataManipulationDto.getContent())
                .collectionIdentifier(dataManipulationDto.getKnowledgeCollection().getCollectionIdentifier())
                .chartType(config.getChatType())
                .vectorModel(dataManipulationDto.getKnowledgeBase().getVectorModel())
                .dimension(DimensionUtil._getDimension(config))
                .baseType(dataManipulationDto.getKnowledgeBase().getBaseType())
                .splitStrategy(dataManipulationDto.getKnowledgeBase().getSplitStrategy())
                .maxParagraphSize(dataManipulationDto.getKnowledgeBase().getMaxParagraphSize())
                .maxOverlapSize(dataManipulationDto.getKnowledgeBase().getMaxOverlapSize())
                .metaDataLinks(dataManipulationDto.getMetaDataLinks())
                .build();
        return embeddingDto;
    }
}
